from my_utils.api_request import *
from my_utils.data_type import QueryResult, QueryResultItem
from dacite import from_dict


def get_chunk_by_fuzzy_search(query, top_k=3):
    search_results = search_papers(query, top_k)
    wraped_search_result = QueryResult(results=[
        from_dict(data_class=QueryResultItem, data=item)
        for item in search_results
    ])
    paper_id_list = []
    for item_result in wraped_search_result.results:
        # 根据paper_id查询论文的第一个chunk
        paper_id_list.append(item_result.entity.paper_id)

    all_paper_data = asyncio.run(fetch_references_and_process(paper_id_list=paper_id_list, top_k=1))

    all_paper_data_str = "\n".join([f"{{\n\"paper_id\": \"{paper['paper_id']}\",\n\"paper_title\": \"{paper['paper_title']}\",\n\"paper_content\": \"{paper['paper_content']}\"\n}}," for paper in all_paper_data])

    return all_paper_data_str


# 根据sql语句查询数据库
tool_get_chunk_by_fuzzy_search = {
    "tool_schema": {
        "type": "function",
        "function": {
            "name": "get_chunk_by_fuzzy_search",
            "strict": True,
            "description": "基于文献查询语句，执行模糊查询，最终返回相关的论文切块",
            "parameters": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "英文查询语句",
                    },
                    "top_k": {
                        "type": "number",
                        "description": "返回的论文切块数量，必须设置为2或3",
                    }
                },
                "required": ["query", "top_k"],
                "additionalProperties": False,
            },
        },
    },
    "tool_def": get_chunk_by_fuzzy_search
}

if __name__ == "__main__":
    sql_result = get_chunk_by_fuzzy_search("机器学习")
    print(sql_result)